import cv2
import numpy as np
import torch
from sklearn import metrics
from sklearn.metrics import precision_recall_fscore_support, mean_absolute_error
import os

from model import RAS
from utils import trans_im, crop_and_flatten

START_ID = 701
END_ID = 1000

if __name__ == "__main__":
    ras = RAS()
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    ras.to(device)

    # 加载模型参数
    ras.load_state_dict(torch.load("data/model/epoch_199_params.pkl", map_location=device))

    auc_list = []
    mae_list = []
    f1_list = []

    for img_id in range(START_ID, END_ID + 1):
        im_path = f"data/train/images/{img_id:04d}.jpg"
        gt_path = f"data/train/ground_truth_mask/{img_id:04d}.png"

        if not os.path.exists(im_path) or not os.path.exists(gt_path):
            print(f"[跳过] 找不到图像或标签：{img_id:04d}")
            continue

        im = cv2.imread(im_path)
        gt = cv2.imread(gt_path)
        if im is None or gt is None:
            print(f"[跳过] 图像读取失败：{img_id:04d}")
            continue

        im_shape = im.shape
        gt = gt[:, :, 0].astype(np.float64) / 255.

        # 预处理
        x = trans_im(im_path)

        # 推理
        y_prob = ras.test(torch.FloatTensor(x).to(device))
        y_prob = y_prob.cpu().numpy()[0, 0, :, :]
        y_prob = cv2.resize(y_prob, (im_shape[1], im_shape[0]), interpolation=cv2.INTER_AREA)

        # 计算指标
        gt_flatten, y_flatten = crop_and_flatten(gt, y_prob)
        if len(gt_flatten) == 0 or len(y_flatten) == 0:
            print(f"[跳过] 空标签或尺寸不匹配：{img_id:04d}")
            continue

        # Ground Truth 二值化
        gt_bin = (gt_flatten > 0.5).astype(np.uint8)
        y_bin = (y_flatten > 0.5).astype(np.uint8)

        # AUC
        try:
            auc = metrics.roc_auc_score(gt_bin, y_flatten)
            auc_list.append(auc)
        except Exception as e:
            print(f"[跳过] 无法计算 AUC：{img_id:04d}, 错误：{e}")

        # MAE
        mae = mean_absolute_error(gt_flatten, y_flatten)
        mae_list.append(mae)

        # F-measure (F1 Score)
        precision, recall, f1, _ = precision_recall_fscore_support(
            gt_bin, y_bin, average='binary', zero_division=0
        )
        f1_list.append(f1)

        print(f"图像 {img_id:04d}: AUC={auc:.4f}, MAE={mae:.4f}, F1={f1:.4f}")

    # 输出平均结果
    if auc_list:
        print(f"\n平均 AUC（共 {len(auc_list)} 张图像）: {np.mean(auc_list):.4f}")
        print(f"平均 MAE: {np.mean(mae_list):.4f}")
        print(f"平均 F-measure (F1 Score): {np.mean(f1_list):.4f}")
    else:
        print("\n没有成功评估任何图像。")
